from typing import Any, Iterable

from zkl_registries.basics import AnyTagKey, AnyTagQuery, AnyTagVal, AnyValue, MutableRegistry, Queries, TagKey, Tags


class SimpleRegistry(MutableRegistry[AnyTagKey, AnyTagVal, AnyValue]):
    def __init__(self,
        keys: Iterable[TagKey],
        records: Iterable[tuple[Tags, AnyValue]] = ()
    ):
        self._keys: tuple[TagKey, ...] = tuple(key for key in keys)
        self._records: list[tuple[dict[TagKey, AnyTagVal], Any]] = []

        for tags, value in records:
            self.register(tags, value)

    @property
    def keys(self) -> tuple[TagKey, ...]:
        return self._keys

    def register(self, tags: Tags, value: AnyValue):
        tags = {key: val for key, val in dict[TagKey, AnyTagVal](tags).items()}
        tags = {key: tags.get(key) for key in self.keys}
        self._records.append((tags, value))

    def match(self, queries: Queries) -> tuple[AnyValue, ...]:
        queries = {key: val for key, val in dict[TagKey, AnyTagQuery](queries).items()}

        records_matching = []
        for tags, value in self._records:
            keys_matching = tuple(key.match(queries.get(key), tags.get(key)) for key in self._keys)
            if all(matching is not None for matching in keys_matching):
                records_matching.append((value, keys_matching))

        sorted_records = sorted(records_matching, key=lambda record: record[1])
        return tuple(record[0] for record in sorted_records)
